/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lvyh.lightframe.tracer.resttemplate.plugin.interceptor;

import com.lvyh.lightframe.tracer.core.SimpleSpan;
import com.lvyh.lightframe.tracer.core.SimpleTracer;
import com.lvyh.lightframe.tracer.core.common.AbstractTracer;
import com.lvyh.lightframe.tracer.core.common.SimpleTracerConfiguration;
import com.lvyh.lightframe.tracer.core.common.TraceContants;
import com.lvyh.lightframe.tracer.core.context.SimpleTraceContext;
import com.lvyh.lightframe.tracer.core.context.SimpleTraceContextHolder;
import com.lvyh.lightframe.tracer.core.samplers.Sampler;
import com.lvyh.lightframe.tracer.resttemplate.plugin.carrier.RestTemplateRequestCarrier;
import io.opentracing.propagation.Format;
import io.opentracing.tag.Tags;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpRequest;
import org.springframework.http.client.ClientHttpRequestExecution;
import org.springframework.http.client.ClientHttpRequestInterceptor;
import org.springframework.http.client.ClientHttpResponse;

import java.io.IOException;
import java.util.List;

public class RestTemplateInterceptor implements ClientHttpRequestInterceptor {
    private static final Logger logger = LoggerFactory.getLogger(RestTemplateInterceptor.class);

    protected AbstractTracer restTemplateTracer;

    public RestTemplateInterceptor(AbstractTracer restTemplateTracer) {
        this.restTemplateTracer = restTemplateTracer;
    }

    @Override
    public ClientHttpResponse intercept(HttpRequest httpRequest, byte[] body, ClientHttpRequestExecution requestExecution) throws IOException {
        ClientHttpResponse httpResponse = null;
        logger.info("[RestTemplateInterceptor] Start intercept request.");

        SimpleSpan simpleSpan = restTemplateTracer.clientSend(httpRequest.getMethod().name());
        appendRestTemplateRequestSpanTags(httpRequest, simpleSpan);
        Throwable t = null;
        try {
            return httpResponse = requestExecution.execute(httpRequest, body);
        } catch (Exception e) {
            logger.error("intercept request error:{}", e);
            t = e;
            throw e;
        } finally {
            SimpleTraceContext simpleTraceContext = SimpleTraceContextHolder.getSimpleTraceContext();
            SimpleSpan currentSpan = simpleTraceContext.getCurrentSpan();
            String resultCode = TraceContants.RESULT_CODE_ERROR;
            if (t != null) {
                currentSpan.setTag(Tags.ERROR.getKey(), t.getMessage());
                simpleSpan.setTag(TraceContants.CURRENT_THREAD_NAME, Thread.currentThread().getName());
            }
            if (httpResponse != null) {
                HttpHeaders headers = httpResponse.getHeaders();
                if (headers != null && headers.get("Content-Length") != null
                        && !headers.get("Content-Length").isEmpty()) {
                    List<String> contentLengthList = headers.get("Content-Length");
                    String len = contentLengthList.get(0);
                    simpleSpan.setTag(TraceContants.RESP_SIZE, Long.valueOf(len));
                }
                simpleSpan.setTag(TraceContants.CURRENT_THREAD_NAME, Thread.currentThread().getName());

                resultCode = String.valueOf(httpResponse.getStatusCode().value());
            }

            Sampler sampler = restTemplateTracer.getSimpleTracer().getSampler();
            //sampler is support &  current span is root span
            if (sampler != null && currentSpan.getParentSpan() == null) {
                currentSpan.getSimpleSpanContext().setSampled(true);
                currentSpan.getSimpleSpanContext().setSampler(sampler);
            }

            restTemplateTracer.clientReceive(resultCode);
        }

    }

    private void appendRestTemplateRequestSpanTags(HttpRequest request, SimpleSpan simpleSpan) {
        if (simpleSpan == null) {
            return;
        }
        simpleSpan.setTag(TraceContants.LOCAL_APP, SimpleTracerConfiguration.getApplicationName());
        simpleSpan.setTag(TraceContants.REQUEST_URL, request.getURI().toString());
        simpleSpan.setTag(TraceContants.METHOD, request.getMethod().name());
        this.injectCarrier(request, simpleSpan);
    }

    public void injectCarrier(HttpRequest request, SimpleSpan currentSpan) {
        SimpleTracer simpleTracer = this.restTemplateTracer.getSimpleTracer();
        simpleTracer.inject(currentSpan.getSimpleSpanContext(), Format.Builtin.HTTP_HEADERS, new RestTemplateRequestCarrier(request));
    }
}
